'''
Author: SlytherinGe
LastEditTime: 2022-02-14 15:51:02
'''
import torch
import torchvision.models as models
import torch.nn as nn
import torch.nn.init as init

resnet_collections = {
    'resnet50': models.resnet50(),
    'resnet101': models.resnet101(),
    'resnet18': models.resnet18(),
    'resnet34': models.resnet34(),
    'resnet152': models.resnet152()
}

class MyResNet(nn.Module):
    def __init__(self, resnet_type, embed_dims):
        super(MyResNet, self).__init__()
        self.layers = []
        resnet = resnet_collections[resnet_type]
        for layer in resnet.children():
            self.layers.append(layer)
        # self.conv1 = self.layers[0]
        self.stem = nn.Conv2d(3, 64, 3, 1, 1)
        self.bn1 = self.layers[1]
        self.relu = self.layers[2]
        self.maxpool = self.layers[3]
        self.layer1 = self.layers[4]
        self.layer2 = self.layers[5]
        self.layer3 = self.layers[6]
        self.layer4 = self.layers[7]
        self.upsample = nn.Upsample(scale_factor=4, mode='bilinear', align_corners=False)
        self.embed_layers = nn.Sequential(
            nn.Conv2d(2048, 512, 3, 1, 1),
            nn.BatchNorm2d(512, affine=True),
            nn.ReLU(),
            nn.Conv2d(512, 256, 3, 1, 1),
            nn.BatchNorm2d(256, affine=True),
            nn.ReLU(),
            nn.Conv2d(256, embed_dims, 1, 1, 0)
        )

    def init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                init.kaiming_normal_(m.weight.detach(), nonlinearity='relu')
                if m.bias is not None:
                    m.bias.detach().fill_(0)
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.detach().fill_(1.0)
                m.bias.detach().zero_()

    def forward(self, x):
        x = self.stem(x)
        # x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.upsample(x)
        out = self.embed_layers(x)

        return out

def create_resnetX_from_pretrained(resnet_type, embed_dims, file_path=None):
    resnet = MyResNet(resnet_type, embed_dims)
    if file_path is not None:
        model_dict = resnet.state_dict()
        print("loading weights...")
        loaded_dict = torch.load(file_path)
        pretrained_dict = {k: v for k, v in loaded_dict.items() if k in model_dict}
        model_dict.update(pretrained_dict)
        resnet.load_state_dict(model_dict)
        print("weights loaded!")
    print('model created!')
    return resnet

